﻿using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;

namespace NanoMvvm.Collections
{
    public class WeakKeyedDictionary<TKey, TValue> : IDictionary<TKey, TValue> where TKey : class
    {
        private readonly Dictionary<WeakReference<TKey>, TValue> internalDictionary = new Dictionary<WeakReference<TKey>, TValue>();
        private int lastCount;
        private long lastGlobalMem;

        public int Count
        {
            get
            {
                Purge();
                return internalDictionary.Count;
            }
        }

        public bool IsReadOnly
        {
            get { return false; }
        }

        public ICollection<TKey> Keys
        {
            get
            {
                return GetKeysEnumerable().ToList();
            }
        }

        public ICollection<TValue> Values
        {
            get
            {
                return GetValueEnumerable().ToList();
            }
        }

        public TValue this[TKey key]
        {
            get
            {
                if (key == null)
                {
                    throw new ArgumentNullException("key");
                }

                Purge();
                return internalDictionary[new WeakReference<TKey>(key)];
            }

            set
            {
                if (key == null)
                {
                    throw new ArgumentNullException("key");
                }

                Purge();
                internalDictionary[new WeakReference<TKey>(key)] = value;
            }
        }

        IEnumerator IEnumerable.GetEnumerator()
        {
            return GetEnumerator();
        }

        void ICollection<KeyValuePair<TKey, TValue>>.Add(KeyValuePair<TKey, TValue> item)
        {
            Add(item.Key, item.Value);
        }

        bool ICollection<KeyValuePair<TKey, TValue>>.Contains(KeyValuePair<TKey, TValue> item)
        {
            TValue value;
            return internalDictionary.TryGetValue(new WeakReference<TKey>(item.Key), out value) && Equals(value, item.Value);
        }

        void ICollection<KeyValuePair<TKey, TValue>>.CopyTo(KeyValuePair<TKey, TValue>[] array, int arrayIndex)
        {
            ((IEnumerable<KeyValuePair<TKey, TValue>>)GetEnumerator()).ToArray().CopyTo(array, arrayIndex);
        }

        bool ICollection<KeyValuePair<TKey, TValue>>.Remove(KeyValuePair<TKey, TValue> item)
        {
            TValue value;
            if (!TryGetValue(item.Key, out value))
            {
                return false;
            }

            return Equals(item.Value, value) && Remove(item.Key);
        }

        public void Purge(bool useMemoryHeuristic = true)
        {
            if (useMemoryHeuristic && !DeadReferencesProbable())
            {
                return;
            }

            var toRemove = internalDictionary.Keys.Where(key => !key.IsTargetAlive).ToList();

            foreach (var entry in toRemove)
            {
                internalDictionary.Remove(entry);
            }
        }

        public IEnumerator<KeyValuePair<TKey, TValue>> GetEnumerator()
        {
            Purge(false);

            return (from entry in internalDictionary
                    select new KeyValuePair<TKey, TValue>(entry.Key.Target, entry.Value))
                    .GetEnumerator();
        }

        public void Clear()
        {
            internalDictionary.Clear();
        }

        public bool ContainsKey(TKey key)
        {
            Purge();
            return internalDictionary.ContainsKey(new WeakReference<TKey>(key));
        }

        public void Add(TKey key, TValue value)
        {
            if (key == null)
            {
                throw new ArgumentNullException("key");
            }
            Purge();
            internalDictionary.Add(new WeakReference<TKey>(key), value);
        }

        public bool Remove(TKey key)
        {
            Purge();
            return internalDictionary.Remove(new WeakReference<TKey>(key));
        }

        public bool TryGetValue(TKey key, out TValue value)
        {
            Purge();
            return internalDictionary.TryGetValue(new WeakReference<TKey>(key), out value);
        }

        private bool DeadReferencesProbable()
        {
            int count = internalDictionary.Count;
            if (lastCount == 0)
            {
                lastCount = count;
                return false;
            }

            long totalMemory = GC.GetTotalMemory(false);
            if (lastGlobalMem == 0L)
            {
                lastGlobalMem = totalMemory;
                return false;
            }

            long memChange = totalMemory - lastGlobalMem;
            int countChange = count - lastCount;

            lastGlobalMem = totalMemory;
            lastCount = count;

            return (memChange < 0L) && (countChange >= 0);
        }

        private IEnumerable<TKey> GetKeysEnumerable()
        {
            Purge(false);

            return internalDictionary.Select(entry => entry.Key.Target);
        }

        private IEnumerable<TValue> GetValueEnumerable()
        {
            Purge(false);

            return internalDictionary.Select(entry => entry.Value);
        }
    }
}